#include <chrono>
#include <cstdlib>
#include <cstring>
#include <iomanip>
#include <iostream>
#include <string>
#include <sys/socket.h>
#include <sys/un.h>
#include <unistd.h>
#include <vector>

#include "bin_resultset_loader.hpp"
#include "cvs_loader.hpp"
#include "mysql_resultset_mock.hpp"
#include "postgresql_resultset_mock.hpp"
#include "uds_client.hpp"

// Benchmark configuration
struct BenchmarkConfig {
    size_t warmup_iterations = 1;
    size_t measurement_iterations = 10;
    size_t payload_size = 0;
};

void ShowUsage(const std::string &program_name) {
    std::cerr << "Usage: " << program_name << " [options] <input_file>\n"
              << "Options:\n"
              << "  -s, --socket <path>    Set Unix domain socket path (default: /tmp/decryption.sock)\n"
              << "  -d, --dbtype TYPE      Database type: mysql or postgresql (default: mysql)\n"
              << "  -w, --warmup <n>       Number of warmup iterations (default: 1)\n"
              << "  -i, --iterations <n>   Number of measurement iterations (default: 10)\n"
              << "  -p, --payload <size>   Use fixed payload size instead of input data (bytes, default: 0)\n"
              << "  -m, --mode <type>      Input file type: csv or resultset (default: csv)\n"
              << "  -h, --help             Show this help message\n";
}

void PrintHex(const std::vector<uint8_t> &response) {
    for (uint8_t byte : response) {
        std::cout << std::hex << std::setw(2) << std::setfill('0') << (int)byte << " ";
    }
    std::cout << std::dec << std::endl; // 恢复为十进制（可选）
}

// 字符串版本
void RunStringBenchmark(UDSClient &client, const std::string &request, const BenchmarkConfig &config) {
    std::vector<uint64_t> latencies;
    latencies.reserve(config.measurement_iterations);

    size_t total_bytes_sent = 0;
    size_t total_bytes_received = 0;

    // Warmup phase
    std::cout << "Running warmup (" << config.warmup_iterations << " iterations)..." << std::endl;
    for (size_t i = 0; i < config.warmup_iterations; ++i) {
        std::string response;
        if (!client.SendRequest(request, response)) {
            std::cerr << "Warmup iteration " << i << " failed" << std::endl;
        }
        // 仅输出 Warmup phase 阶段的测试内容
        std::cout << "Received response from server (" << response.size() << " bytes):\n";
        std::cout << response << std::endl;
    }

    // Measurement phase
    std::cout << "Running benchmark (" << config.measurement_iterations << " iterations)..." << std::endl;
    for (size_t i = 0; i < config.measurement_iterations; ++i) {
        std::string response;
        if (client.SendRequest(request, response)) {
            latencies.push_back(client.GetLastDurationNs());
            total_bytes_sent += client.GetLastBytesSent();
            total_bytes_received += client.GetLastBytesReceived();
        } else {
            std::cerr << "Measurement iteration " << i << " failed" << std::endl;
        }
    }

    // Calculate statistics
    if (latencies.empty()) {
        std::cerr << "Error: No successful measurements recorded" << std::endl;
        return;
    }

    uint64_t total_ns = 0;
    uint64_t min_ns = latencies[0];
    uint64_t max_ns = latencies[0];

    for (auto ns : latencies) {
        total_ns += ns;
        if (ns < min_ns)
            min_ns = ns;
        if (ns > max_ns)
            max_ns = ns;
    }

    double avg_ns = static_cast<double>(total_ns) / latencies.size();
    double avg_ms = avg_ns / 1e6;
    double min_ms = min_ns / 1e6;
    double max_ms = max_ns / 1e6;

    double throughput_mbps = (total_bytes_sent + total_bytes_received) * 8 / (total_ns / 1e9) / 1e6;

    // Print results
    std::cout << "\nBenchmark Results:\n";
    std::cout << "----------------------------------------\n";
    std::cout << "Iterations:        " << latencies.size() << "\n";
    std::cout << "Payload size:     ~" << request.size() / 1024 << " KB\n";
    std::cout << "Average latency:   " << std::fixed << std::setprecision(2) << avg_ms << " ms\n";
    std::cout << "Minimum latency:   " << std::fixed << std::setprecision(2) << min_ms << " ms\n";
    std::cout << "Maximum latency:   " << std::fixed << std::setprecision(2) << max_ms << " ms\n";
    std::cout << "Total throughput:  " << std::fixed << std::setprecision(2) << throughput_mbps << " Mbps\n";
    std::cout << "Total data sent:    " << total_bytes_sent / 1024 << " KB\n";
    std::cout << "Total data received: " << total_bytes_received / 1024 << " KB\n";
    std::cout << "----------------------------------------\n";
}

// 二进制数据版本
void RunBinaryBenchmark(UDSClient &client, const std::vector<uint8_t> &request, const BenchmarkConfig &config, const std::string dbtype) {
    std::vector<uint64_t> latencies;
    latencies.reserve(config.measurement_iterations);

    size_t total_bytes_sent = 0;
    size_t total_bytes_received = 0;

    // Warmup phase
    std::cout << "Running warmup (" << config.warmup_iterations << " iterations)..." << std::endl;
    for (size_t i = 0; i < config.warmup_iterations; ++i) {
        std::vector<uint8_t> response;
        if (!client.SendRequest(request, response)) {
            std::cerr << "Warmup iteration " << i << " failed" << std::endl;
        }

        // 仅输出 Warmup phase 阶段的测试内容
        std::cout << "Received response from server (" << response.size() << " bytes):\n";
        std::cout << "Test parsing result set type: " << dbtype << std::endl;

        if (dbtype == "mysql") {
            MySQLResultSetParser parser;
            try {
                parser.parse(response);

                std::cout << "Packet Number: " << static_cast<int>(parser.getPacketNumber()) << std::endl;
                std::cout << "Schema: " << parser.getSchema() << std::endl;
                std::cout << "Table: " << parser.getTable() << std::endl;
                std::cout << "Column Count: " << parser.getColumnCount() << std::endl;

                for (const auto &col : parser.getColumns()) {
                    std::cout << "Column: " << col.name << " (" << col.type << ")" << std::endl;
                }
            } catch (const std::exception &e) {
                std::cerr << "Error: " << e.what() << std::endl;
            }
        } else if (dbtype == "postgresql") {
            PostgresqlResultSetParser parser;
            try {
                parser.parse(response);

                std::cout << "Table Oid: " << parser.getFirstTableOid() << std::endl;
                std::cout << "Column Count: " << parser.getColumnCount() << std::endl;

                for (const auto &col : parser.getColumns()) {
                    std::cout << "Column: " << col.name << " (" << col.type << ")" << std::endl;
                }
            } catch (const std::exception &e) {
                std::cerr << "Error: " << e.what() << std::endl;
            }
            // PrintHex(response);
        } else {
            std::cerr << "Invalid database type: " << dbtype << " (must be mysql or postgresql)\n";
            return;
        }
    }

    // Measurement phase
    std::cout << "Running benchmark (" << config.measurement_iterations << " iterations)..." << std::endl;
    for (size_t i = 0; i < config.measurement_iterations; ++i) {
        std::vector<uint8_t> response;
        if (client.SendRequest(request, response)) {
            latencies.push_back(client.GetLastDurationNs());
            total_bytes_sent += client.GetLastBytesSent();
            total_bytes_received += client.GetLastBytesReceived();
        } else {
            std::cerr << "Measurement iteration " << i << " failed" << std::endl;
        }
    }

    // Calculate statistics
    if (latencies.empty()) {
        std::cerr << "Error: No successful measurements recorded" << std::endl;
        return;
    }

    uint64_t total_ns = 0;
    uint64_t min_ns = latencies[0];
    uint64_t max_ns = latencies[0];

    for (auto ns : latencies) {
        total_ns += ns;
        if (ns < min_ns)
            min_ns = ns;
        if (ns > max_ns)
            max_ns = ns;
    }

    double avg_ns = static_cast<double>(total_ns) / latencies.size();
    double avg_ms = avg_ns / 1e6;
    double min_ms = min_ns / 1e6;
    double max_ms = max_ns / 1e6;

    double throughput_mbps = (total_bytes_sent + total_bytes_received) * 8 / (total_ns / 1e9) / 1e6;

    // Print results
    std::cout << "\nBenchmark Results:\n";
    std::cout << "----------------------------------------\n";
    std::cout << "Iterations:        " << latencies.size() << "\n";
    std::cout << "Payload size:     ~" << request.size() / 1024 << " KB\n";
    std::cout << "Average latency:   " << std::fixed << std::setprecision(2) << avg_ms << " ms\n";
    std::cout << "Minimum latency:   " << std::fixed << std::setprecision(2) << min_ms << " ms\n";
    std::cout << "Maximum latency:   " << std::fixed << std::setprecision(2) << max_ms << " ms\n";
    std::cout << "Total throughput:  " << std::fixed << std::setprecision(2) << throughput_mbps << " Mbps\n";
    std::cout << "Total data sent:    " << total_bytes_sent / 1024 << " KB\n";
    std::cout << "Total data received: " << total_bytes_received / 1024 << " KB\n";
    std::cout << "----------------------------------------\n";
}

int main(int argc, char *argv[]) {
    // Default parameters
    std::string socket_path = "/tmp/decryption.sock";
    std::string csv_file;
    BenchmarkConfig config;
    std::string input_mode = "csv"; // 默认 CSV 模式
    std::string db_type = "mysql";  // 默认数据库类型

    // Parse command line arguments
    for (int i = 1; i < argc; ++i) {
        std::string arg = argv[i];
        if (arg == "-h" || arg == "--help") {
            ShowUsage(argv[0]);
            return 0;
        } else if ((arg == "-s" || arg == "--socket") && i + 1 < argc) {
            socket_path = argv[++i];
        } else if ((arg == "-d" || arg == "--dbtype") && i + 1 < argc) {
            db_type = argv[++i];
            if (db_type != "mysql" && db_type != "postgresql") {
                std::cerr << "Invalid database type: " << db_type << " (must be mysql or postgresql)\n";
                return 1;
            }
        } else if ((arg == "-w" || arg == "--warmup") && i + 1 < argc) {
            config.warmup_iterations = std::stoul(argv[++i]);
        } else if ((arg == "-i" || arg == "--iterations") && i + 1 < argc) {
            config.measurement_iterations = std::stoul(argv[++i]);
        } else if ((arg == "-p" || arg == "--payload") && i + 1 < argc) {
            config.payload_size = std::stoul(argv[++i]);
        } else if ((arg == "-m" || arg == "--mode") && i + 1 < argc) {
            input_mode = argv[++i];
            if (input_mode != "csv" && input_mode != "resultset") {
                std::cerr << "Invalid mode: " << input_mode << " (must be csv or resultset)\n";
                return 1;
            }
        } else if (arg[0] != '-') {
            csv_file = arg;
        } else {
            std::cerr << "Unknown option: " << arg << std::endl;
            ShowUsage(argv[0]);
            return 1;
        }
    }

    try {
        // Create UDS client
        std::cout << "Connecting to UDS server at " << socket_path << std::endl;
        UDSClient uds_client(socket_path);

        if (config.payload_size > 0) {
            // Generate synthetic payload
            std::string request(config.payload_size, 'A');
            std::cout << "Using synthetic payload of " << config.payload_size << " bytes\n";
            RunStringBenchmark(uds_client, request, config);
        } else {
            if (input_mode == "csv") {
                // Load CSV data
                std::cout << "Loading CSV file: " << csv_file << std::endl;
                CvsResult mysqlData = loadCSVResult(csv_file);
                std::string requestData = mysqlData.column_names + "\n";
                for (const auto &row : mysqlData.rows) {
                    for (const auto &field : row) {
                        requestData += field + ",";
                    }
                    requestData += "\n";
                }
                RunStringBenchmark(uds_client, requestData, config);
            } else {
                // Load binary resultset
                std::cout << "Loading binary ResultSet file: " << csv_file << std::endl;
                std::vector<uint8_t> binaryData = BinaryFileReader::loadBinaryFile(csv_file);
                // BinaryFileReader::printHex(binaryData); // 打印十六进制验证
                RunBinaryBenchmark(uds_client, binaryData, config, db_type);
            }
        }

        return 0;
    } catch (const std::exception &e) {
        std::cerr << "\nError: " << e.what() << std::endl;
        return 1;
    } catch (...) {
        std::cerr << "\nUnknown error occurred!" << std::endl;
        return 1;
    }
}